import torch
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans
from pandas.core.frame import DataFrame
import pandas as pd
import numpy as np
import argparse
import yaml
import matplotlib.pyplot as plt
import numba

# 用TSNE进行数据降维并展示聚类结果

parser = argparse.ArgumentParser()
parser.add_argument('--feat_file', type=str, default='./input')
parser.add_argument('--cls_file', type=str, default='./input')

args = parser.parse_args()

text_feat = torch.load(args.feat_file)

with open(args.cls_file, 'r') as f:
    cls_data = yaml.load(f, Loader=yaml.Loader)

feat_arr = torch.stack(list(text_feat.values()), dim=0).numpy()
cls_arr = np.array(list(cls_data.values()), dtype=int)

n_cls = 20
feat_dim = feat_arr[0].shape[0]

center_list = np.zeros((n_cls, feat_dim))
for i in range(n_cls):
    feat_i = feat_arr[cls_arr - 1 == i]
    feat_mean = np.mean(feat_i, axis=0)
    feat_dist = np.sum(np.square(feat_i - feat_mean[None, :]) / feat_dim, axis=-1)

    dist_std_f = np.std(feat_dist) * 1.5
    dist_mean = np.mean(feat_dist)

    plt.figure()
    plt.hist(feat_dist, 500, facecolor='blue', alpha=0.5)
    ymin, ymax = plt.gca().get_ylim()
    plt.vlines(dist_mean - dist_std_f, ymin, ymax, linestyles="dashed", colors="b")
    plt.vlines(dist_mean + dist_std_f, ymin, ymax, linestyles="dashed", colors="#ecb116")
    plt.show()

    keep_idx = np.abs(feat_dist - dist_mean) < dist_std_f
    center_list[i] = np.mean(feat_i[keep_idx, :], axis=0)

np.save('center_list.npy', center_list)